# Copyright 2025 the LlamaFactory team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import subprocess
import sys
from copy import deepcopy
from functools import partial


USAGE = (
    "-" * 70
    + "\n"
    + "| Usage:                                                             |\n"
    + "|   llamafactory-cli api -h: launch an OpenAI-style API server       |\n"
    + "|   llamafactory-cli chat -h: launch a chat interface in CLI         |\n"
    + "|   llamafactory-cli eval -h: evaluate models                        |\n"
    + "|   llamafactory-cli export -h: merge LoRA adapters and export model |\n"
    + "|   llamafactory-cli train -h: train models                          |\n"
    + "|   llamafactory-cli webchat -h: launch a chat interface in Web UI   |\n"
    + "|   llamafactory-cli webui: launch LlamaBoard                        |\n"
    + "|   llamafactory-cli version: show version info                      |\n"
    + "-" * 70
)


def main():
    from . import launcher
    from .api.app import run_api
    # chat模块
    from .chat.chat_model import run_chat
    # 评估模块
    from .eval.evaluator import run_eval
    # 用于日志记录。
    from .extras import logging
    # 从 .extras.env 导入 VERSION（版本号）和 print_env（打印环境信息）。
    from .extras.env import VERSION, print_env
    # 查找可用端口，GPU数量，是否使用 Ray 进行分布式训练
    from .extras.misc import find_available_port, get_device_count, is_env_enabled, use_ray
    # 分别用于导出模型和运行训练任务。
    from .train.tuner import export_model, run_exp
    # 从 .webui.interface 导入 run_web_demo 和 run_web_ui，分别用于运行 Web Demo 和完整的 Web UI。
    from .webui.interface import run_web_demo, run_web_ui

    logger = logging.get_logger(__name__)

    WELCOME = (
        "-" * 58
        + "\n"
        + f"| Welcome to LLaMA Factory, version {VERSION}"
        + " " * (21 - len(VERSION))
        + "|\n|"
        + " " * 56
        + "|\n"
        + "| Project page: https://github.com/hiyouga/LLaMA-Factory |\n"
        + "-" * 58
    )

    # 命令映射表
    COMMAND_MAP = {
        "api": run_api,
        "chat": run_chat,
        "env": print_env,
        "eval": run_eval,
        "export": export_model,
        "train": run_exp,
        "webchat": run_web_demo,
        "webui": run_web_ui,
        "version": partial(print, WELCOME),
        "help": partial(print, USAGE),
    }

    # 解析用户输入的第一个命令参数，若无则默认为 help
    command = sys.argv.pop(1) if len(sys.argv) > 1 else "help"

    # 如果是 train 命令，并且启用了 FORCE_TORCHRUN 或者有多个 GPU 但未使用 Ray，则使用 torchrun 启动分布式训练
    if command == "train" and (is_env_enabled("FORCE_TORCHRUN") or (get_device_count() > 1 and not use_ray())):
        # launch distributed training
        # 获取节点数环境变量，默认为 1
        nnodes = os.getenv("NNODES", "1")
        # 获取当前节点的排名，默认为 0
        node_rank = os.getenv("NODE_RANK", "0")
        # 获取每个节点使用的 GPU 数量，默认为当前设备数量
        nproc_per_node = os.getenv("NPROC_PER_NODE", str(get_device_count()))
        # 获取主节点地址，默认为本地回环地址
        master_addr = os.getenv("MASTER_ADDR", "127.0.0.1")
        # 获取主节点通信端口，若未设置则自动查找一个可用端口
        master_port = os.getenv("MASTER_PORT", str(find_available_port()))
        # 输出分布式训练初始化信息（只在 rank0 节点输出）
        logger.info_rank0(f"Initializing {nproc_per_node} distributed tasks at: {master_addr}:{master_port}")
        if int(nnodes) > 1:
            # 如果节点数大于1，表示启用多节点训练
            print(f"Multi-node training enabled: num nodes: {nnodes}, node rank: {node_rank}")

        # 拷贝当前环境变量，供后续子进程使用
        env = deepcopy(os.environ)

        # 如果启用了优化 Torch 的选项，则设置相应的环境变量以提升性能
        if is_env_enabled("OPTIM_TORCH", "1"):
            # optimize DDP, see https://zhuanlan.zhihu.com/p/671834539
            # 设置 CUDA 内存分配策略和 NCCL 参数优化 DDP 性能
            env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
            env["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"

        # NOTE: DO NOT USE shell=True to avoid security risk
        # 使用 subprocess.run 执行 torchrun 命令启动分布式训练
        """
        --nnodes: 节点总数（机器数），例如在多台服务器上训练时使用。总共多少台机器
        --node_rank: 当前节点的编号（从0开始）。
        --nproc_per_node: 每个节点使用的 GPU 数量。
        --master_addr: 主节点地址，默认是本地 127.0.0.1。
        --master_port: 主节点通信端口，自动选择一个可用端口。
        {file_name}: 要执行的 Python 文件路径，这里是 launcher.py。
        {args}: 剩下的命令行参数，比如模型配置、数据集路径等。
        举例子：
        torchrun --nnodes 1 --node_rank 0 --nproc_per_node 2 \
            --master_addr 127.0.0.1 --master_port 29500 \
            launcher.py train --config data/config.yaml
        然后点击launcher进入
        """
        process = subprocess.run(
            (
                "torchrun --nnodes {nnodes} --node_rank {node_rank} --nproc_per_node {nproc_per_node} "
                "--master_addr {master_addr} --master_port {master_port} {file_name} {args}"
            )
            .format(
                nnodes=nnodes,
                node_rank=node_rank,
                nproc_per_node=nproc_per_node,
                master_addr=master_addr,
                master_port=master_port,
                file_name=launcher.__file__,
                args=" ".join(sys.argv[1:]),
            )
            .split(),
            env=env,
            check=True,
        )
        # 子进程结束后退出并返回其状态码
        sys.exit(process.returncode)
    elif command in COMMAND_MAP:
        # 如果命令合法，则调用对应函数执行
        COMMAND_MAP[command]()
    else:
        # 如果命令非法，输出错误提示和帮助信息
        print(f"Unknown command: {command}.\n{USAGE}")


if __name__ == "__main__":
    from multiprocessing import freeze_support

    freeze_support()
    main()
